package com.mrx.springboot.demo.utils;

import lombok.SneakyThrows;
import org.apache.commons.codec.binary.Hex;
import org.slf4j.MDC;
import org.slf4j.helpers.MessageFormatter;

import java.net.Inet4Address;
import java.net.InetAddress;
import java.net.NetworkInterface;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author Mr.X
 * @since 2025/2/12 21:00
 */
public class TracerUtils {

    public static final String TRACE_ID = "traceId";

    private static final AtomicInteger THREAD_COUNT = new AtomicInteger();

    private static final AtomicInteger SEQUENCE = new AtomicInteger(1000);

    private static final Inet4Address ADDRESS = getAddress();

    public static void startTracer() {
        MDC.put(TRACE_ID, getTraceId());
    }

    public static void stopTracer() {
        MDC.remove(TRACE_ID);
    }

    public static Thread startTracer(Runnable r) {
        String traceId = getTraceId();
        return new Thread(() -> {
            MDC.put(TRACE_ID, traceId);
            r.run();
        });
    }

    public static Thread startTracer(Thread thread) {
        String traceId = getTraceId();
        return new Thread(thread.getThreadGroup(), () -> {
            MDC.put(TRACE_ID, traceId);
            ((Runnable) thread).run();
        }, thread.getName());
    }

    public static ExecutorService newFixedTracerPool(int threads, String namePrefix) {
        return newFixedThreadPool(threads, TracerUtils::startTracer, namePrefix);
    }

    @SneakyThrows
    public static void safeSleep(long sleepTime, TimeUnit timeUnit) {
        timeUnit.sleep(sleepTime);
    }

    public static ExecutorService newFixedThreadPool(int nThreads, ThreadFactory threadFactory, String namePrefix) {
        TimeUnit unit = TimeUnit.MILLISECONDS;
        LinkedBlockingQueue<Runnable> workQueue = new LinkedBlockingQueue<>();
        ThreadFactory decoratedThreadFactory = it -> decorateThread(threadFactory.newThread(it), namePrefix);
        return new XThreadPoolExecutor(nThreads, nThreads, 0L, unit, workQueue, decoratedThreadFactory);
    }

    @SneakyThrows
    public static String parseTraceId(String traceId) {
        List<String> result = new ArrayList<>();
        for (Integer index : Arrays.asList(8, 13, 4, 5)) {
            int idx = Math.min(traceId.length(), index);
            result.add(traceId.substring(0, idx));
            traceId = traceId.substring(idx);
        }
        InetAddress ip = InetAddress.getByAddress(Hex.decodeHex(result.getFirst()));
        String time = result.get(1);
        String seq = result.get(2);
        String pid = result.get(3);
        return format("ip: {}, time: {}, seq: {}, pid: {}", ip.getHostAddress(), time, seq, pid);
    }

    public static String format(String template, Object... args) {
        return MessageFormatter.basicArrayFormat(template, args);
    }

    private static Thread decorateThread(Thread thread, String namePrefix) {
        thread.setName(namePrefix + "-" + THREAD_COUNT.getAndIncrement());
        return thread;
    }

    private static final class XThreadPoolExecutor extends ThreadPoolExecutor {

        public XThreadPoolExecutor(
                int corePoolSize,
                int maximumPoolSize,
                long keepAliveTime,
                TimeUnit unit,
                BlockingQueue<Runnable> workQueue,
                ThreadFactory threadFactory
        ) {
            super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue,
                    threadFactory);
        }

        @Override
        protected void afterExecute(Runnable r, Throwable t) {
            super.afterExecute(r, t);
        }

    }

    private static String getTraceId() {
        return Optional.ofNullable(MDC.get(TRACE_ID)).orElseGet(() -> {
            try {
                return getSofaTraceId();
            } catch (Throwable t) {
                return UUID.randomUUID().toString();
            }
        });
    }

    private static String getSofaTraceId() {
        StringBuilder sb = new StringBuilder();
        for (String part : ADDRESS.getHostAddress().split("\\.")) {
            sb.append(String.format("%02x", Integer.parseInt(part)));
        }
        sb.append(System.currentTimeMillis());
        int seq = SEQUENCE.getAndIncrement();
        if (seq <= 9000) {
            sb.append(seq);
        } else {
            SEQUENCE.set(1000);
            sb.append(1000);
        }
        long pid = ProcessHandle.current().pid();
        sb.append(pid > 100000 ? pid / 100000 : pid);
        return sb.toString();
    }

    @SneakyThrows
    private static Inet4Address getAddress() {
        return NetworkInterface.networkInterfaces().filter(TracerUtils::filterLoopback)
                .map(NetworkInterface::getInetAddresses).map(Collections::list)
                .flatMap(Collection::stream)
                .filter(Inet4Address.class::isInstance)
                .findFirst().map(Inet4Address.class::cast)
                .orElseGet(TracerUtils::getLocalHost);
    }

    @SneakyThrows
    private static Inet4Address getLocalHost() {
        return (Inet4Address) InetAddress.getByName("127.0.0.1");
    }

    @SneakyThrows
    private static boolean filterLoopback(NetworkInterface networkInterface) {
        return !networkInterface.isLoopback();
    }

}
